import numpy as np
import math

class TwoDofKinematicModel:
    """
    车辆二自由度运动学模型类
    状态变量: [x, y, theta, v]
    控制输入: [加速度, 前轮转角]
    """
    
    def __init__(self, x=0.0, y=0.0, theta=0.0, v=0.0, L=2.5, dt=0.1):
        """
        初始化车辆状态
        
        Parameters:
        -----------
        x : float, 初始x坐标 (m)
        y : float, 初始y坐标 (m) 
        theta : float, 初始航向角 (rad)
        v : float, 初始速度 (m/s)
        L : float, 轴距 (m)
        dt : float, 时间步长 (s)
        """
        self.state = np.array([x, y, theta, v], dtype=np.float32)
        self.L = L  # 轴距
        self.dt = dt  # 时间步长
        self.history = []  # 历史状态记录
        self.history.append(self.state.copy())
    
    def update(self, acceleration, delta):
        """
        更新车辆状态（欧拉积分）
        
        Parameters:
        -----------
        acceleration : float, 加速度 (m/s²)
        delta : float, 前轮转角 (rad)
        
        Returns:
        --------
        new_state : ndarray, 更新后的状态 [x, y, theta, v]
        """
        x, y, theta, v = self.state
        
        # 限制前轮转角在物理范围内
        delta = np.clip(delta, -np.pi/3, np.pi/3)
        
        # 运动学方程
        dx = v * np.cos(theta)
        dy = v * np.sin(theta)
        dtheta = (v / self.L) * np.tan(delta)
        dv = acceleration
        
        # 欧拉积分更新状态
        new_x = x + dx * self.dt
        new_y = y + dy * self.dt
        new_theta = theta + dtheta * self.dt
        new_v = v + dv * self.dt
        
        # 限制速度不为负（简单处理）
        new_v = max(new_v, 0)
        
        # 归一化角度到 [-pi, pi]
        new_theta = self.normalize_angle(new_theta)
        
        # 更新状态
        self.state = np.array([new_x, new_y, new_theta, new_v])
        self.history.append(self.state.copy())
        
        return self.state
    
    def get_state(self):
        """
        获取当前状态信息
        
        Returns:
        --------
        dict: 包含所有状态信息的字典
        """
        x, y, theta, v = self.state
        
        return {
            'position': (x, y),
            'x': x,
            'y': y,
            'theta': theta,
            'heading_deg': np.degrees(theta),
            'velocity': v,
            'state_vector': self.state.copy()
        }
    
    def get_history(self):
        """
        获取历史轨迹
        
        Returns:
        --------
        ndarray: N×4的数组，每行是[x, y, theta, v]
        """
        return np.array(self.history)
    
    def reset(self, x=0.0, y=0.0, theta=0.0, v=0.0):
        """
        重置车辆状态
        
        Parameters:
        -----------
        x : float, 新的x坐标
        y : float, 新的y坐标
        theta : float, 新的航向角
        v : float, 新的速度
        """
        self.state = np.array([x, y, theta, v], dtype=np.float32)
        self.history = [self.state.copy()]
    
    def predict_trajectory(self, acceleration, delta, steps=10):
        """
        预测未来轨迹
        
        Parameters:
        -----------
        acceleration : float, 加速度
        delta : float, 前轮转角
        steps : int, 预测步数
        
        Returns:
        --------
        trajectory : ndarray, 预测轨迹 [x, y, theta, v] × steps
        """
        # 保存当前状态
        current_state = self.state.copy()
        
        trajectory = []
        for _ in range(steps):
            self.update(acceleration, delta)
            trajectory.append(self.state.copy())
        
        # 恢复原始状态
        self.state = current_state
        
        return np.array(trajectory)
    
    @staticmethod
    def normalize_angle(angle):
        """
        将角度归一化到 [-pi, pi] 范围内
        
        Parameters:
        -----------
        angle : float, 输入角度 (rad)
        
        Returns:
        --------
        normalized_angle : float, 归一化后的角度
        """
        while angle > np.pi:
            angle -= 2 * np.pi
        while angle < -np.pi:
            angle += 2 * np.pi
        return angle
    
    def __str__(self):
        """返回状态信息的字符串表示"""
        state_info = self.get_state()
        return (f"Vehicle State: \n"
                f"  Position: ({state_info['x']:.2f}, {state_info['y']:.2f}) m\n"
                f"  Heading: {state_info['theta']:.3f} rad ({state_info['heading_deg']:.1f}°)\n"
                f"  Velocity: {state_info['velocity']:.2f} m/s")


# 使用示例
if __name__ == "__main__":
    # 创建车辆模型
    vehicle = TwoDofKinematicModel(x=0, y=0, theta=0, v=5.0, L=2.5, dt=0.1)
    
    print("初始状态:")
    print(vehicle)
    print()
    
    # 模拟控制输入：轻微转向
    acceleration = 0.0  # 保持匀速
    delta = np.radians(10)  # 10度转向
    
    # 更新10步
    for i in range(10):
        new_state = vehicle.update(acceleration, delta)
        if i % 2 == 0:  # 每2步打印一次
            print(f"Step {i+1}:")
            print(f"  Position: ({new_state[0]:.2f}, {new_state[1]:.2f})")
            print(f"  Heading: {np.degrees(new_state[2]):.1f}°")
            print(f"  Velocity: {new_state[3]:.2f} m/s")
            print()
    
    # 获取最终状态信息
    final_info = vehicle.get_state()
    print("最终状态信息:")
    for key, value in final_info.items():
        if key != 'state_vector':
            print(f"  {key}: {value}")
    
    # 获取历史轨迹
    history = vehicle.get_history()
    print(f"\n轨迹点数: {len(history)}")